{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "460b5e31",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import Descriptors\n",
    "\n",
    "from chemprop.utils import make_mol\n",
    "\n",
    "smiles_list = [\n",
    "    \"[H][H]\",\n",
    "    \"C\",\n",
    "    \"CN\",\n",
    "    \"CN\",\n",
    "    \"CC\",\n",
    "    \"[CH2:3]=[N+:1]([H:4])[H:2]\",\n",
    "    \"CCCC\",\n",
    "    \"CO\",\n",
    "    \"CC#N\",\n",
    "    \"C1NN1\",\n",
    "    \"c1cc[n-]c1\",\n",
    "]\n",
    "mols = [make_mol(smi, keep_h=True, add_h=False, reorder_atoms=True) for smi in smiles_list]\n",
    "columns = [\"smiles\", \"mol_y1\", \"mol_y2\", \"atom_y1\", \"atom_y2\", \"bond_y1\", \"bond_y2\", \"weight\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d28aa677",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_d = []\n",
    "V_fds = []\n",
    "E_fds = []\n",
    "\n",
    "for i, mol in enumerate(mols):\n",
    "    x_d1 = int(any(atom.GetAtomicNum() == 6 for atom in mol.GetAtoms()))  # Has carbon\n",
    "    x_d2 = int(any(atom.GetAtomicNum() == 7 for atom in mol.GetAtoms()))  # Has nitrogen\n",
    "    v_fd = np.array(\n",
    "        [[int(atom.GetAtomicNum() == 6), int(atom.GetAtomicNum() == 7)] for atom in mol.GetAtoms()]\n",
    "    )  # Is carbon, Is nitrogen\n",
    "    e_fd = np.array(\n",
    "        [\n",
    "            [\n",
    "                int(bond.GetBondType() == Chem.rdchem.BondType.SINGLE),\n",
    "                int(bond.GetBondType() == Chem.rdchem.BondType.DOUBLE),\n",
    "            ]\n",
    "            for bond in mol.GetBonds()\n",
    "        ]\n",
    "    ).reshape(\n",
    "        mol.GetNumBonds(), 2\n",
    "    )  # Is single bond, Is double bond\n",
    "\n",
    "    X_d.append([x_d1, x_d2])\n",
    "    V_fds.append(v_fd)\n",
    "    E_fds.append(e_fd)\n",
    "\n",
    "np.savez(\"atom_features_descriptors.npz\", *V_fds)\n",
    "np.savez(\"bond_features_descriptors.npz\", *E_fds)\n",
    "np.savez(\"descriptors.npz\", np.array(X_d))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3c6db394",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([[0, 0],\n",
       "  [1, 0],\n",
       "  [1, 1],\n",
       "  [1, 1],\n",
       "  [1, 0],\n",
       "  [1, 1],\n",
       "  [1, 0],\n",
       "  [1, 0],\n",
       "  [1, 1],\n",
       "  [1, 1],\n",
       "  [1, 1]],\n",
       " [array([[0, 0],\n",
       "         [0, 0]]),\n",
       "  array([[1, 0]]),\n",
       "  array([[1, 0],\n",
       "         [0, 1]]),\n",
       "  array([[1, 0],\n",
       "         [0, 1]]),\n",
       "  array([[1, 0],\n",
       "         [1, 0]]),\n",
       "  array([[0, 1],\n",
       "         [0, 0],\n",
       "         [1, 0],\n",
       "         [0, 0]]),\n",
       "  array([[1, 0],\n",
       "         [1, 0],\n",
       "         [1, 0],\n",
       "         [1, 0]]),\n",
       "  array([[1, 0],\n",
       "         [0, 0]]),\n",
       "  array([[1, 0],\n",
       "         [1, 0],\n",
       "         [0, 1]]),\n",
       "  array([[1, 0],\n",
       "         [0, 1],\n",
       "         [0, 1]]),\n",
       "  array([[1, 0],\n",
       "         [1, 0],\n",
       "         [1, 0],\n",
       "         [0, 1],\n",
       "         [1, 0]])],\n",
       " [array([[1, 0]]),\n",
       "  array([], shape=(0, 2), dtype=float64),\n",
       "  array([[1, 0]]),\n",
       "  array([[1, 0]]),\n",
       "  array([[1, 0]]),\n",
       "  array([[0, 1],\n",
       "         [1, 0],\n",
       "         [1, 0]]),\n",
       "  array([[1, 0],\n",
       "         [1, 0],\n",
       "         [1, 0]]),\n",
       "  array([[1, 0]]),\n",
       "  array([[1, 0],\n",
       "         [0, 0]]),\n",
       "  array([[1, 0],\n",
       "         [1, 0],\n",
       "         [1, 0]]),\n",
       "  array([[0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]])])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_d, V_fds, E_fds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5220c148",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>mol_y1</th>\n",
       "      <th>mol_y2</th>\n",
       "      <th>atom_y1</th>\n",
       "      <th>atom_y2</th>\n",
       "      <th>bond_y1</th>\n",
       "      <th>bond_y2</th>\n",
       "      <th>weight</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[H][H]</td>\n",
       "      <td>2.016</td>\n",
       "      <td>2.0</td>\n",
       "      <td>[1, 1]</td>\n",
       "      <td>[1.008, 1.008]</td>\n",
       "      <td>[2]</td>\n",
       "      <td>[-2]</td>\n",
       "      <td>0.090909</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>C</td>\n",
       "      <td>16.043</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[6]</td>\n",
       "      <td>[12.011]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0.181818</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CN</td>\n",
       "      <td>31.058</td>\n",
       "      <td>2.0</td>\n",
       "      <td>[6, 7]</td>\n",
       "      <td>[12.011, 14.007]</td>\n",
       "      <td>[13]</td>\n",
       "      <td>[-13]</td>\n",
       "      <td>0.272727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CN</td>\n",
       "      <td>31.058</td>\n",
       "      <td>NaN</td>\n",
       "      <td>[6, 7]</td>\n",
       "      <td>[None, 14.007]</td>\n",
       "      <td>[13]</td>\n",
       "      <td>[None]</td>\n",
       "      <td>0.363636</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CC</td>\n",
       "      <td>30.070</td>\n",
       "      <td>2.0</td>\n",
       "      <td>[6, 6]</td>\n",
       "      <td>[12.011, 12.011]</td>\n",
       "      <td>[12]</td>\n",
       "      <td>[-12]</td>\n",
       "      <td>0.454545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[CH2:3]=[N+:1]([H:4])[H:2]</td>\n",
       "      <td>30.050</td>\n",
       "      <td>4.0</td>\n",
       "      <td>[7, 1, 6, 1]</td>\n",
       "      <td>[14.007, 1.008, 12.011, 1.008]</td>\n",
       "      <td>[13, 8, 8]</td>\n",
       "      <td>[-13, -8, -8]</td>\n",
       "      <td>0.545455</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>CCCC</td>\n",
       "      <td>58.124</td>\n",
       "      <td>4.0</td>\n",
       "      <td>[6, 6, 6, 6]</td>\n",
       "      <td>[12.011, 12.011, 12.011, 12.011]</td>\n",
       "      <td>[12, 12, 12]</td>\n",
       "      <td>[-12, -12, -12]</td>\n",
       "      <td>0.636364</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>CO</td>\n",
       "      <td>32.042</td>\n",
       "      <td>2.0</td>\n",
       "      <td>[6, 8]</td>\n",
       "      <td>[12.011, 15.999]</td>\n",
       "      <td>[14]</td>\n",
       "      <td>[-14]</td>\n",
       "      <td>0.727273</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>CC#N</td>\n",
       "      <td>41.053</td>\n",
       "      <td>3.0</td>\n",
       "      <td>[6, 6, 7]</td>\n",
       "      <td>[12.011, 12.011, 14.007]</td>\n",
       "      <td>[12, 13]</td>\n",
       "      <td>[-12, -13]</td>\n",
       "      <td>0.818182</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>C1NN1</td>\n",
       "      <td>44.057</td>\n",
       "      <td>3.0</td>\n",
       "      <td>[6, 7, 7]</td>\n",
       "      <td>[12.011, 14.007, 14.007]</td>\n",
       "      <td>[13, 14, 13]</td>\n",
       "      <td>[-13, -14, -13]</td>\n",
       "      <td>0.909091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>c1cc[n-]c1</td>\n",
       "      <td>66.083</td>\n",
       "      <td>5.0</td>\n",
       "      <td>[6, 6, 6, 7, 6]</td>\n",
       "      <td>[12.011, 12.011, 12.011, 14.007, 12.011]</td>\n",
       "      <td>[12, 12, 13, 13, 12]</td>\n",
       "      <td>[-12, -12, -13, -13, -12]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        smiles  mol_y1  mol_y2          atom_y1  \\\n",
       "0                       [H][H]   2.016     2.0           [1, 1]   \n",
       "1                            C  16.043     1.0              [6]   \n",
       "2                           CN  31.058     2.0           [6, 7]   \n",
       "3                           CN  31.058     NaN           [6, 7]   \n",
       "4                           CC  30.070     2.0           [6, 6]   \n",
       "5   [CH2:3]=[N+:1]([H:4])[H:2]  30.050     4.0     [7, 1, 6, 1]   \n",
       "6                         CCCC  58.124     4.0     [6, 6, 6, 6]   \n",
       "7                           CO  32.042     2.0           [6, 8]   \n",
       "8                         CC#N  41.053     3.0        [6, 6, 7]   \n",
       "9                        C1NN1  44.057     3.0        [6, 7, 7]   \n",
       "10                  c1cc[n-]c1  66.083     5.0  [6, 6, 6, 7, 6]   \n",
       "\n",
       "                                     atom_y2               bond_y1  \\\n",
       "0                             [1.008, 1.008]                   [2]   \n",
       "1                                   [12.011]                    []   \n",
       "2                           [12.011, 14.007]                  [13]   \n",
       "3                             [None, 14.007]                  [13]   \n",
       "4                           [12.011, 12.011]                  [12]   \n",
       "5             [14.007, 1.008, 12.011, 1.008]            [13, 8, 8]   \n",
       "6           [12.011, 12.011, 12.011, 12.011]          [12, 12, 12]   \n",
       "7                           [12.011, 15.999]                  [14]   \n",
       "8                   [12.011, 12.011, 14.007]              [12, 13]   \n",
       "9                   [12.011, 14.007, 14.007]          [13, 14, 13]   \n",
       "10  [12.011, 12.011, 12.011, 14.007, 12.011]  [12, 12, 13, 13, 12]   \n",
       "\n",
       "                      bond_y2    weight  \n",
       "0                        [-2]  0.090909  \n",
       "1                          []  0.181818  \n",
       "2                       [-13]  0.272727  \n",
       "3                      [None]  0.363636  \n",
       "4                       [-12]  0.454545  \n",
       "5               [-13, -8, -8]  0.545455  \n",
       "6             [-12, -12, -12]  0.636364  \n",
       "7                       [-14]  0.727273  \n",
       "8                  [-12, -13]  0.818182  \n",
       "9             [-13, -14, -13]  0.909091  \n",
       "10  [-12, -12, -13, -13, -12]  1.000000  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = []\n",
    "for i, (mol, smiles) in enumerate(zip(mols, smiles_list)):\n",
    "    # Molecular properties\n",
    "    mol_y1 = Descriptors.MolWt(mol)  # Molecular weight\n",
    "    mol_y2 = mol.GetNumAtoms()  # Number of atoms\n",
    "\n",
    "    # Atom properties\n",
    "    atom_y1 = [atom.GetAtomicNum() for atom in mol.GetAtoms()]  # Atomic numbers\n",
    "    atom_y2 = [atom.GetMass() for atom in mol.GetAtoms()]  # Atomic masses\n",
    "\n",
    "    # Bond properties\n",
    "    bond_y1 = [\n",
    "        sum((bond.GetBeginAtom().GetAtomicNum(), bond.GetEndAtom().GetAtomicNum()))\n",
    "        for bond in mol.GetBonds()\n",
    "    ]  # Sum of atomic numbers of bonded atoms\n",
    "    bond_y2 = [-x for x in bond_y1]  # Negative of bond_y1\n",
    "\n",
    "    weight = (i + 1) / len(smiles_list)\n",
    "\n",
    "    if i == 3:\n",
    "        mol_y2 = None\n",
    "        atom_y2[0] = None\n",
    "        bond_y2[0] = None\n",
    "\n",
    "    data.append([smiles, mol_y1, mol_y2, atom_y1, atom_y2, bond_y1, bond_y2, weight])\n",
    "\n",
    "df = pd.DataFrame(data, columns=columns)\n",
    "df.to_csv(\"regression.csv\", index=False)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "93ab3980",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>mol_y1</th>\n",
       "      <th>mol_y2</th>\n",
       "      <th>atom_y1</th>\n",
       "      <th>atom_y2</th>\n",
       "      <th>bond_y1</th>\n",
       "      <th>bond_y2</th>\n",
       "      <th>weight</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[H][H]</td>\n",
       "      <td>&lt;2.016</td>\n",
       "      <td>2</td>\n",
       "      <td>[&lt;1, &gt;1]</td>\n",
       "      <td>[&gt;1.008, 1.008]</td>\n",
       "      <td>[&lt;2]</td>\n",
       "      <td>[&gt;-2]</td>\n",
       "      <td>0.090909</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>C</td>\n",
       "      <td>&lt;16.043</td>\n",
       "      <td>1</td>\n",
       "      <td>[6]</td>\n",
       "      <td>[&lt;12.011]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0.181818</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CN</td>\n",
       "      <td>&lt;31.058</td>\n",
       "      <td>2</td>\n",
       "      <td>[6, 7]</td>\n",
       "      <td>[&gt;12.011, &gt;14.007]</td>\n",
       "      <td>[&lt;13]</td>\n",
       "      <td>[&lt;-13]</td>\n",
       "      <td>0.272727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CN</td>\n",
       "      <td>&gt;31.058</td>\n",
       "      <td>None</td>\n",
       "      <td>[&lt;6, 7]</td>\n",
       "      <td>[None, 14.007]</td>\n",
       "      <td>[&gt;13]</td>\n",
       "      <td>[None]</td>\n",
       "      <td>0.363636</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CC</td>\n",
       "      <td>&lt;30.07</td>\n",
       "      <td>&gt;2</td>\n",
       "      <td>[&gt;6, &lt;6]</td>\n",
       "      <td>[12.011, &lt;12.011]</td>\n",
       "      <td>[&gt;12]</td>\n",
       "      <td>[&lt;-12]</td>\n",
       "      <td>0.454545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[CH2:3]=[N+:1]([H:4])[H:2]</td>\n",
       "      <td>30.05</td>\n",
       "      <td>&gt;4</td>\n",
       "      <td>[&gt;7, 1, 6, 1]</td>\n",
       "      <td>[&gt;14.007, &lt;1.008, &gt;12.011, 1.008]</td>\n",
       "      <td>[&lt;13, &gt;8, &gt;8]</td>\n",
       "      <td>[-13, -8, &lt;-8]</td>\n",
       "      <td>0.545455</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>CCCC</td>\n",
       "      <td>58.123999999999995</td>\n",
       "      <td>4</td>\n",
       "      <td>[6, &lt;6, &gt;6, &lt;6]</td>\n",
       "      <td>[&lt;12.011, &gt;12.011, &gt;12.011, 12.011]</td>\n",
       "      <td>[12, &gt;12, &gt;12]</td>\n",
       "      <td>[-12, &gt;-12, &gt;-12]</td>\n",
       "      <td>0.636364</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>CO</td>\n",
       "      <td>32.042</td>\n",
       "      <td>2</td>\n",
       "      <td>[&lt;6, &lt;8]</td>\n",
       "      <td>[&lt;12.011, &lt;15.999]</td>\n",
       "      <td>[&gt;14]</td>\n",
       "      <td>[-14]</td>\n",
       "      <td>0.727273</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>CC#N</td>\n",
       "      <td>41.053</td>\n",
       "      <td>&gt;3</td>\n",
       "      <td>[&gt;6, 6, 7]</td>\n",
       "      <td>[&gt;12.011, &gt;12.011, &gt;14.007]</td>\n",
       "      <td>[12, 13]</td>\n",
       "      <td>[&lt;-12, &gt;-13]</td>\n",
       "      <td>0.818182</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>C1NN1</td>\n",
       "      <td>&lt;44.057</td>\n",
       "      <td>3</td>\n",
       "      <td>[6, 7, 7]</td>\n",
       "      <td>[&lt;12.011, 14.007, &lt;14.007]</td>\n",
       "      <td>[&lt;13, 14, &gt;13]</td>\n",
       "      <td>[&lt;-13, &lt;-14, -13]</td>\n",
       "      <td>0.909091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>c1cc[n-]c1</td>\n",
       "      <td>66.08299999999998</td>\n",
       "      <td>&lt;5</td>\n",
       "      <td>[&gt;6, 6, 6, &gt;7, 6]</td>\n",
       "      <td>[&lt;12.011, 12.011, 12.011, &lt;14.007, &gt;12.011]</td>\n",
       "      <td>[12, &lt;12, &lt;13, &lt;13, &gt;12]</td>\n",
       "      <td>[&lt;-12, &lt;-12, -13, &gt;-13, &gt;-12]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        smiles              mol_y1 mol_y2            atom_y1  \\\n",
       "0                       [H][H]              <2.016      2           [<1, >1]   \n",
       "1                            C             <16.043      1                [6]   \n",
       "2                           CN             <31.058      2             [6, 7]   \n",
       "3                           CN             >31.058   None            [<6, 7]   \n",
       "4                           CC              <30.07     >2           [>6, <6]   \n",
       "5   [CH2:3]=[N+:1]([H:4])[H:2]               30.05     >4      [>7, 1, 6, 1]   \n",
       "6                         CCCC  58.123999999999995      4    [6, <6, >6, <6]   \n",
       "7                           CO              32.042      2           [<6, <8]   \n",
       "8                         CC#N              41.053     >3         [>6, 6, 7]   \n",
       "9                        C1NN1             <44.057      3          [6, 7, 7]   \n",
       "10                  c1cc[n-]c1   66.08299999999998     <5  [>6, 6, 6, >7, 6]   \n",
       "\n",
       "                                        atom_y2                   bond_y1  \\\n",
       "0                               [>1.008, 1.008]                      [<2]   \n",
       "1                                     [<12.011]                        []   \n",
       "2                            [>12.011, >14.007]                     [<13]   \n",
       "3                                [None, 14.007]                     [>13]   \n",
       "4                             [12.011, <12.011]                     [>12]   \n",
       "5             [>14.007, <1.008, >12.011, 1.008]             [<13, >8, >8]   \n",
       "6           [<12.011, >12.011, >12.011, 12.011]            [12, >12, >12]   \n",
       "7                            [<12.011, <15.999]                     [>14]   \n",
       "8                   [>12.011, >12.011, >14.007]                  [12, 13]   \n",
       "9                    [<12.011, 14.007, <14.007]            [<13, 14, >13]   \n",
       "10  [<12.011, 12.011, 12.011, <14.007, >12.011]  [12, <12, <13, <13, >12]   \n",
       "\n",
       "                          bond_y2    weight  \n",
       "0                           [>-2]  0.090909  \n",
       "1                              []  0.181818  \n",
       "2                          [<-13]  0.272727  \n",
       "3                          [None]  0.363636  \n",
       "4                          [<-12]  0.454545  \n",
       "5                  [-13, -8, <-8]  0.545455  \n",
       "6               [-12, >-12, >-12]  0.636364  \n",
       "7                           [-14]  0.727273  \n",
       "8                    [<-12, >-13]  0.818182  \n",
       "9               [<-13, <-14, -13]  0.909091  \n",
       "10  [<-12, <-12, -13, >-13, >-12]  1.000000  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bounds = [\"<\", \"\", \">\"]\n",
    "\n",
    "data = []\n",
    "for i, (mol, smiles) in enumerate(zip(mols, smiles_list)):\n",
    "    # Molecular properties\n",
    "    mol_y1 = Descriptors.MolWt(mol)  # Molecular weight\n",
    "    mol_y2 = mol.GetNumAtoms()  # Number of atoms\n",
    "\n",
    "    # Atom properties\n",
    "    atom_y1 = [atom.GetAtomicNum() for atom in mol.GetAtoms()]  # Atomic numbers\n",
    "    atom_y2 = [atom.GetMass() for atom in mol.GetAtoms()]  # Atomic masses\n",
    "\n",
    "    # Bond properties\n",
    "    bond_y1 = [\n",
    "        sum((bond.GetBeginAtom().GetAtomicNum(), bond.GetEndAtom().GetAtomicNum()))\n",
    "        for bond in mol.GetBonds()\n",
    "    ]  # Sum of atomic numbers of bonded atoms\n",
    "    bond_y2 = [-x for x in bond_y1]  # Negative of bond_y1\n",
    "\n",
    "    weight = (i + 1) / len(smiles_list)\n",
    "\n",
    "    mol_y1 = bounds[np.random.randint(0, 3)] + str(mol_y1)\n",
    "    mol_y2 = bounds[np.random.randint(0, 3)] + str(mol_y2)\n",
    "    atom_y1 = [bounds[np.random.randint(0, 3)] + str(x) for x in atom_y1]\n",
    "    atom_y2 = [bounds[np.random.randint(0, 3)] + str(x) for x in atom_y2]\n",
    "    bond_y1 = [bounds[np.random.randint(0, 3)] + str(x) for x in bond_y1]\n",
    "    bond_y2 = [bounds[np.random.randint(0, 3)] + str(x) for x in bond_y2]\n",
    "\n",
    "    if i == 3:\n",
    "        mol_y2 = None\n",
    "        atom_y2[0] = None\n",
    "        bond_y2[0] = None\n",
    "\n",
    "    data.append([smiles, mol_y1, mol_y2, atom_y1, atom_y2, bond_y1, bond_y2, weight])\n",
    "\n",
    "df = pd.DataFrame(data, columns=columns)\n",
    "df.to_csv(\"bounded.csv\", index=False)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "75e879fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>mol_y1</th>\n",
       "      <th>mol_y2</th>\n",
       "      <th>atom_y1</th>\n",
       "      <th>atom_y2</th>\n",
       "      <th>bond_y1</th>\n",
       "      <th>bond_y2</th>\n",
       "      <th>weight</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[H][H]</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>0.090909</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>C</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0.181818</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CN</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[1, 0]</td>\n",
       "      <td>[0, 1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>0.272727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CN</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>[1, 0]</td>\n",
       "      <td>[None, 1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[None]</td>\n",
       "      <td>0.363636</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CC</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[1, 1]</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>0.454545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[CH2:3]=[N+:1]([H:4])[H:2]</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[0, 0, 1, 0]</td>\n",
       "      <td>[1, 0, 0, 0]</td>\n",
       "      <td>[0, 1, 1]</td>\n",
       "      <td>[1, 0, 0]</td>\n",
       "      <td>0.545455</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>CCCC</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[1, 1, 1, 1]</td>\n",
       "      <td>[0, 0, 0, 0]</td>\n",
       "      <td>[1, 1, 1]</td>\n",
       "      <td>[0, 0, 0]</td>\n",
       "      <td>0.636364</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>CO</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[1, 0]</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>0.727273</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>CC#N</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[1, 1, 0]</td>\n",
       "      <td>[0, 0, 1]</td>\n",
       "      <td>[1, 0]</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>0.818182</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>C1NN1</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[1, 0, 0]</td>\n",
       "      <td>[0, 1, 1]</td>\n",
       "      <td>[1, 1, 1]</td>\n",
       "      <td>[0, 0, 0]</td>\n",
       "      <td>0.909091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>c1cc[n-]c1</td>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[1, 1, 1, 0, 1]</td>\n",
       "      <td>[0, 0, 0, 1, 0]</td>\n",
       "      <td>[0, 0, 0, 0, 0]</td>\n",
       "      <td>[0, 0, 0, 0, 0]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        smiles  mol_y1  mol_y2          atom_y1  \\\n",
       "0                       [H][H]       0     0.0           [0, 0]   \n",
       "1                            C       1     0.0              [1]   \n",
       "2                           CN       1     1.0           [1, 0]   \n",
       "3                           CN       1     NaN           [1, 0]   \n",
       "4                           CC       1     0.0           [1, 1]   \n",
       "5   [CH2:3]=[N+:1]([H:4])[H:2]       1     1.0     [0, 0, 1, 0]   \n",
       "6                         CCCC       1     0.0     [1, 1, 1, 1]   \n",
       "7                           CO       1     0.0           [1, 0]   \n",
       "8                         CC#N       1     1.0        [1, 1, 0]   \n",
       "9                        C1NN1       1     1.0        [1, 0, 0]   \n",
       "10                  c1cc[n-]c1       1     1.0  [1, 1, 1, 0, 1]   \n",
       "\n",
       "            atom_y2          bond_y1          bond_y2    weight  \n",
       "0            [0, 0]              [1]              [0]  0.090909  \n",
       "1               [0]               []               []  0.181818  \n",
       "2            [0, 1]              [1]              [0]  0.272727  \n",
       "3         [None, 1]              [1]           [None]  0.363636  \n",
       "4            [0, 0]              [1]              [0]  0.454545  \n",
       "5      [1, 0, 0, 0]        [0, 1, 1]        [1, 0, 0]  0.545455  \n",
       "6      [0, 0, 0, 0]        [1, 1, 1]        [0, 0, 0]  0.636364  \n",
       "7            [0, 0]              [1]              [0]  0.727273  \n",
       "8         [0, 0, 1]           [1, 0]           [0, 0]  0.818182  \n",
       "9         [0, 1, 1]        [1, 1, 1]        [0, 0, 0]  0.909091  \n",
       "10  [0, 0, 0, 1, 0]  [0, 0, 0, 0, 0]  [0, 0, 0, 0, 0]  1.000000  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = []\n",
    "for i, (mol, smiles) in enumerate(zip(mols, smiles_list)):\n",
    "    # Molecule-level features\n",
    "    mol_y1 = int(any(atom.GetAtomicNum() == 6 for atom in mol.GetAtoms()))  # Has carbon\n",
    "    mol_y2 = int(any(atom.GetAtomicNum() == 7 for atom in mol.GetAtoms()))  # Has nitrogen\n",
    "\n",
    "    # Atom-level features\n",
    "    atom_y1 = [int(atom.GetAtomicNum() == 6) for atom in mol.GetAtoms()]  # Is carbon\n",
    "    atom_y2 = [int(atom.GetAtomicNum() == 7) for atom in mol.GetAtoms()]  # Is nitrogen\n",
    "\n",
    "    # Bond-level features\n",
    "    bond_y1 = [\n",
    "        int(bond.GetBondType() == Chem.rdchem.BondType.SINGLE) for bond in mol.GetBonds()\n",
    "    ]  # Is single bond\n",
    "    bond_y2 = [\n",
    "        int(bond.GetBondType() == Chem.rdchem.BondType.DOUBLE) for bond in mol.GetBonds()\n",
    "    ]  # Is double bond\n",
    "\n",
    "    weight = (i + 1) / len(smiles_list)\n",
    "\n",
    "    if i == 3:\n",
    "        mol_y2 = None\n",
    "        atom_y2[0] = None\n",
    "        bond_y2[0] = None\n",
    "\n",
    "    data.append([smiles, mol_y1, mol_y2, atom_y1, atom_y2, bond_y1, bond_y2, weight])\n",
    "\n",
    "df = pd.DataFrame(data, columns=columns)\n",
    "df.to_csv(\"classification.csv\", index=False)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f897d852",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>mol_y1</th>\n",
       "      <th>mol_y2</th>\n",
       "      <th>atom_y1</th>\n",
       "      <th>atom_y2</th>\n",
       "      <th>bond_y1</th>\n",
       "      <th>bond_y2</th>\n",
       "      <th>weight</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[H][H]</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[2, 2]</td>\n",
       "      <td>[2, 2]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>0.090909</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>C</td>\n",
       "      <td>1</td>\n",
       "      <td>2.0</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0.181818</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CN</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[0, 1]</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>0.272727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CN</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>[0, 1]</td>\n",
       "      <td>[None, 0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[None]</td>\n",
       "      <td>0.363636</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CC</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[2]</td>\n",
       "      <td>0.454545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[CH2:3]=[N+:1]([H:4])[H:2]</td>\n",
       "      <td>1</td>\n",
       "      <td>2.0</td>\n",
       "      <td>[1, 2, 0, 2]</td>\n",
       "      <td>[1, 2, 1, 2]</td>\n",
       "      <td>[2, 0, 0]</td>\n",
       "      <td>[1, 0, 0]</td>\n",
       "      <td>0.545455</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>CCCC</td>\n",
       "      <td>2</td>\n",
       "      <td>2.0</td>\n",
       "      <td>[0, 0, 0, 0]</td>\n",
       "      <td>[0, 0, 0, 0]</td>\n",
       "      <td>[0, 0, 0]</td>\n",
       "      <td>[2, 2, 2]</td>\n",
       "      <td>0.636364</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>CO</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>[0, 2]</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>0.727273</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>CC#N</td>\n",
       "      <td>2</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[0, 0, 1]</td>\n",
       "      <td>[0, 2, 2]</td>\n",
       "      <td>[0, 1]</td>\n",
       "      <td>[2, 1]</td>\n",
       "      <td>0.818182</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>C1NN1</td>\n",
       "      <td>2</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[0, 1, 1]</td>\n",
       "      <td>[0, 0, 0]</td>\n",
       "      <td>[0, 0, 0]</td>\n",
       "      <td>[1, 0, 1]</td>\n",
       "      <td>0.909091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>c1cc[n-]c1</td>\n",
       "      <td>2</td>\n",
       "      <td>2.0</td>\n",
       "      <td>[0, 0, 0, 1, 0]</td>\n",
       "      <td>[1, 1, 1, 1, 1]</td>\n",
       "      <td>[2, 2, 2, 2, 2]</td>\n",
       "      <td>[2, 2, 1, 1, 2]</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        smiles  mol_y1  mol_y2          atom_y1  \\\n",
       "0                       [H][H]       0     0.0           [2, 2]   \n",
       "1                            C       1     2.0              [0]   \n",
       "2                           CN       1     0.0           [0, 1]   \n",
       "3                           CN       1     NaN           [0, 1]   \n",
       "4                           CC       1     0.0           [0, 0]   \n",
       "5   [CH2:3]=[N+:1]([H:4])[H:2]       1     2.0     [1, 2, 0, 2]   \n",
       "6                         CCCC       2     2.0     [0, 0, 0, 0]   \n",
       "7                           CO       1     0.0           [0, 2]   \n",
       "8                         CC#N       2     1.0        [0, 0, 1]   \n",
       "9                        C1NN1       2     1.0        [0, 1, 1]   \n",
       "10                  c1cc[n-]c1       2     2.0  [0, 0, 0, 1, 0]   \n",
       "\n",
       "            atom_y2          bond_y1          bond_y2    weight  \n",
       "0            [2, 2]              [0]              [0]  0.090909  \n",
       "1               [0]               []               []  0.181818  \n",
       "2            [0, 0]              [0]              [1]  0.272727  \n",
       "3         [None, 0]              [0]           [None]  0.363636  \n",
       "4            [0, 0]              [0]              [2]  0.454545  \n",
       "5      [1, 2, 1, 2]        [2, 0, 0]        [1, 0, 0]  0.545455  \n",
       "6      [0, 0, 0, 0]        [0, 0, 0]        [2, 2, 2]  0.636364  \n",
       "7            [0, 0]              [0]              [1]  0.727273  \n",
       "8         [0, 2, 2]           [0, 1]           [2, 1]  0.818182  \n",
       "9         [0, 0, 0]        [0, 0, 0]        [1, 0, 1]  0.909091  \n",
       "10  [1, 1, 1, 1, 1]  [2, 2, 2, 2, 2]  [2, 2, 1, 1, 2]  1.000000  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol_classes1 = lambda x: 0 if x < 10 else 1 if x <= 40 else 2\n",
    "mol_classes2 = {2: 0, 3: 1}\n",
    "atom_classes1 = {6: 0, 7: 1}\n",
    "atom_classes2 = {Chem.rdchem.HybridizationType.SP3: 0, Chem.rdchem.HybridizationType.SP2: 1}\n",
    "bond_classes1 = {Chem.rdchem.BondType.SINGLE: 0, Chem.rdchem.BondType.TRIPLE: 1}\n",
    "bond_classes2 = {0: 0, 1: 1}\n",
    "\n",
    "data = []\n",
    "for i, (mol, smiles) in enumerate(zip(mols, smiles_list)):\n",
    "    # Molecule-level features\n",
    "    mol_y1 = mol_classes1(Descriptors.MolWt(mol))  # Molecular weight class\n",
    "    mol_y2 = mol_classes2.get(mol.GetNumAtoms(), 2)  # Number of atoms class\n",
    "\n",
    "    # Atom-level features\n",
    "    atom_y1 = [\n",
    "        atom_classes1.get(atom.GetAtomicNum(), 2) for atom in mol.GetAtoms()\n",
    "    ]  # Atomic number class\n",
    "    atom_y2 = [\n",
    "        atom_classes2.get(atom.GetHybridization(), 2) for atom in mol.GetAtoms()\n",
    "    ]  # Hybridization class\n",
    "\n",
    "    # Bond-level features\n",
    "    bond_y1 = [\n",
    "        bond_classes1.get(bond.GetBondType(), 2) for bond in mol.GetBonds()\n",
    "    ]  # Bond type class\n",
    "    bond_y2 = [\n",
    "        bond_classes2.get(\n",
    "            sum([a.GetAtomicNum() == 6 for a in [bond.GetBeginAtom(), bond.GetEndAtom()]]), 2\n",
    "        )\n",
    "        for bond in mol.GetBonds()\n",
    "    ]  # Bond type class\n",
    "\n",
    "    weight = (i + 1) / len(smiles_list)\n",
    "\n",
    "    if i == 3:\n",
    "        mol_y2 = None\n",
    "        atom_y2[0] = None\n",
    "        bond_y2[0] = None\n",
    "\n",
    "    data.append([smiles, mol_y1, mol_y2, atom_y1, atom_y2, bond_y1, bond_y2, weight])\n",
    "\n",
    "df = pd.DataFrame(data, columns=columns)\n",
    "df.to_csv(\"multiclass.csv\", index=False)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7e68ae06",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>mol_y</th>\n",
       "      <th>atom_y1</th>\n",
       "      <th>atom_y2</th>\n",
       "      <th>bond_y1</th>\n",
       "      <th>bond_y2</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[H][H]</td>\n",
       "      <td>0</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[1.008, 1.008]</td>\n",
       "      <td>[2]</td>\n",
       "      <td>[2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>C</td>\n",
       "      <td>0</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[12.011]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CN</td>\n",
       "      <td>0</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[12.011, 14.007]</td>\n",
       "      <td>[13]</td>\n",
       "      <td>[2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CN</td>\n",
       "      <td>0</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[12.011, 14.007]</td>\n",
       "      <td>[13]</td>\n",
       "      <td>[2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CC</td>\n",
       "      <td>0</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[12.011, 12.011]</td>\n",
       "      <td>[12]</td>\n",
       "      <td>[2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[CH2:3]=[N+:1]([H:4])[H:2]</td>\n",
       "      <td>1</td>\n",
       "      <td>[1, 0, 0, 0]</td>\n",
       "      <td>[14.007, 1.008, 12.011, 1.008]</td>\n",
       "      <td>[13, 8, 8]</td>\n",
       "      <td>[4, 2, 2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>CCCC</td>\n",
       "      <td>0</td>\n",
       "      <td>[0, 0, 0, 0]</td>\n",
       "      <td>[12.011, 12.011, 12.011, 12.011]</td>\n",
       "      <td>[12, 12, 12]</td>\n",
       "      <td>[2, 2, 2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>CO</td>\n",
       "      <td>0</td>\n",
       "      <td>[0, 0]</td>\n",
       "      <td>[12.011, 15.999]</td>\n",
       "      <td>[14]</td>\n",
       "      <td>[2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>CC#N</td>\n",
       "      <td>0</td>\n",
       "      <td>[0, 0, 0]</td>\n",
       "      <td>[12.011, 12.011, 14.007]</td>\n",
       "      <td>[12, 13]</td>\n",
       "      <td>[2, 6]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>C1NN1</td>\n",
       "      <td>0</td>\n",
       "      <td>[0, 0, 0]</td>\n",
       "      <td>[12.011, 14.007, 14.007]</td>\n",
       "      <td>[13, 14, 13]</td>\n",
       "      <td>[2, 2, 2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>c1cc[n-]c1</td>\n",
       "      <td>-1</td>\n",
       "      <td>[0, 0, 0, -1, 0]</td>\n",
       "      <td>[12.011, 12.011, 12.011, 14.007, 12.011]</td>\n",
       "      <td>[12, 12, 13, 13, 12]</td>\n",
       "      <td>[3, 3, 3, 3, 3]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        smiles  mol_y           atom_y1  \\\n",
       "0                       [H][H]      0            [0, 0]   \n",
       "1                            C      0               [0]   \n",
       "2                           CN      0            [0, 0]   \n",
       "3                           CN      0            [0, 0]   \n",
       "4                           CC      0            [0, 0]   \n",
       "5   [CH2:3]=[N+:1]([H:4])[H:2]      1      [1, 0, 0, 0]   \n",
       "6                         CCCC      0      [0, 0, 0, 0]   \n",
       "7                           CO      0            [0, 0]   \n",
       "8                         CC#N      0         [0, 0, 0]   \n",
       "9                        C1NN1      0         [0, 0, 0]   \n",
       "10                  c1cc[n-]c1     -1  [0, 0, 0, -1, 0]   \n",
       "\n",
       "                                     atom_y2               bond_y1  \\\n",
       "0                             [1.008, 1.008]                   [2]   \n",
       "1                                   [12.011]                    []   \n",
       "2                           [12.011, 14.007]                  [13]   \n",
       "3                           [12.011, 14.007]                  [13]   \n",
       "4                           [12.011, 12.011]                  [12]   \n",
       "5             [14.007, 1.008, 12.011, 1.008]            [13, 8, 8]   \n",
       "6           [12.011, 12.011, 12.011, 12.011]          [12, 12, 12]   \n",
       "7                           [12.011, 15.999]                  [14]   \n",
       "8                   [12.011, 12.011, 14.007]              [12, 13]   \n",
       "9                   [12.011, 14.007, 14.007]          [13, 14, 13]   \n",
       "10  [12.011, 12.011, 12.011, 14.007, 12.011]  [12, 12, 13, 13, 12]   \n",
       "\n",
       "            bond_y2  \n",
       "0               [2]  \n",
       "1                []  \n",
       "2               [2]  \n",
       "3               [2]  \n",
       "4               [2]  \n",
       "5         [4, 2, 2]  \n",
       "6         [2, 2, 2]  \n",
       "7               [2]  \n",
       "8            [2, 6]  \n",
       "9         [2, 2, 2]  \n",
       "10  [3, 3, 3, 3, 3]  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>atom_y1_constraint</th>\n",
       "      <th>atom_y2_constraint</th>\n",
       "      <th>bond_y2_constraint</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>2.016</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>12.011</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>26.018</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>26.018</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>24.022</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1</td>\n",
       "      <td>28.034</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>0</td>\n",
       "      <td>48.044</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>0</td>\n",
       "      <td>28.010</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>0</td>\n",
       "      <td>38.029</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>0</td>\n",
       "      <td>40.025</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>-1</td>\n",
       "      <td>62.051</td>\n",
       "      <td>15</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    atom_y1_constraint  atom_y2_constraint  bond_y2_constraint\n",
       "0                    0               2.016                   2\n",
       "1                    0              12.011                   0\n",
       "2                    0              26.018                   2\n",
       "3                    0              26.018                   2\n",
       "4                    0              24.022                   2\n",
       "5                    1              28.034                   8\n",
       "6                    0              48.044                   6\n",
       "7                    0              28.010                   2\n",
       "8                    0              38.029                   8\n",
       "9                    0              40.025                   6\n",
       "10                  -1              62.051                  15"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = []\n",
    "constraints = []\n",
    "for i, (mol, smiles) in enumerate(zip(mols, smiles_list)):\n",
    "    # Molecular properties\n",
    "    mol_y = Chem.GetFormalCharge(mol)  # Formal charge\n",
    "\n",
    "    # Atom properties\n",
    "    atom_y1 = [atom.GetFormalCharge() for atom in mol.GetAtoms()]  # Atomic charges\n",
    "    atom_y2 = [atom.GetMass() for atom in mol.GetAtoms()]  # Atomic masses\n",
    "\n",
    "    # Bond properties\n",
    "    bond_y1 = [\n",
    "        sum((bond.GetBeginAtom().GetAtomicNum(), bond.GetEndAtom().GetAtomicNum()))\n",
    "        for bond in mol.GetBonds()\n",
    "    ]  # Sum of atomic numbers of bonded atoms\n",
    "    n_electrons = {\n",
    "        Chem.rdchem.BondType.SINGLE: 2,\n",
    "        Chem.rdchem.BondType.AROMATIC: 3,\n",
    "        Chem.rdchem.BondType.DOUBLE: 4,\n",
    "        Chem.rdchem.BondType.TRIPLE: 6,\n",
    "    }\n",
    "    bond_y2 = [n_electrons[bond.GetBondType()] for bond in mol.GetBonds()]\n",
    "\n",
    "    data.append([smiles, mol_y, atom_y1, atom_y2, bond_y1, bond_y2])\n",
    "    constraints.append([mol_y, sum(atom_y2), sum(bond_y2)])\n",
    "\n",
    "df = pd.DataFrame(data, columns=[\"smiles\", \"mol_y\", \"atom_y1\", \"atom_y2\", \"bond_y1\", \"bond_y2\"])\n",
    "df.to_csv(\"constrained_regression.csv\", index=False)\n",
    "display(df)\n",
    "\n",
    "df_constraints = pd.DataFrame(\n",
    "    constraints, columns=[\"atom_y1_constraint\", \"atom_y2_constraint\", \"bond_y2_constraint\"]\n",
    ")\n",
    "df_constraints.to_csv(\"constrained_regression_constraints.csv\", index=False)\n",
    "df_constraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a45e8bc8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chemprop",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
